import math
import os
from typing import List

import torch
from torch import nn
from torch.nn.modules.transformer import _get_clones

from lib.models.layers.head import build_box_head

from lib.models.dutrack.itpn import fast_itpn_base_3324_patch16_224
from lib.utils.box_ops import box_xyxy_to_cxcywh


class DUTrack(nn.Module):
    """ This is the base class for MMTrack """

    def __init__(self, transformer, box_head, aux_loss=False, head_type="CORNER", token_len=1):
        """ Initializes the model.
        Parameters:
            transformer: torch module of the transformer architecture.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
        """
        super().__init__()
        self.backbone = transformer
        self.box_head = box_head

        self.aux_loss = aux_loss
        self.head_type = head_type
        if head_type == "CORNER" or head_type == "CENTER":
            self.feat_sz_s = int(box_head.feat_sz)
            self.feat_len_s = int(box_head.feat_sz ** 2)

        if self.aux_loss:
            self.box_head = _get_clones(self.box_head, 6)

        self.track_query = None
        self.token_len = token_len
        # 新增: Token Bank FIFO队列
        self.token_bank = []
        self.token_bank_max_len = 4  # FIFO队列最大长度

        # 新增: 融合历史token的Learnable Query
        embed_dim = transformer.embed_dim
        self.fusion_query = nn.Parameter(torch.zeros(1, token_len, embed_dim))
        nn.init.normal_(self.fusion_query, std=0.02)

        # 新增: 融合历史token的注意力机制
        self.fusion_attn = nn.MultiheadAttention(
            embed_dim=embed_dim,
            num_heads=8,
            batch_first=True
        )

        # 新增: 与当前搜索区域交互的注意力(可选)
        self.search_rebuild_attn = nn.MultiheadAttention(
            embed_dim=embed_dim,
            num_heads=8,
            batch_first=True
        )

        # 新增: 标准化层和MLP
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.fusion_mlp = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            nn.GELU(),
            nn.Linear(embed_dim * 4, embed_dim)
        )

    def forward(self, template: torch.Tensor,
                search: torch.Tensor,
                descript,
                gamma_scale=1.0,  # 添加gamma_scale参数
                ):
        assert isinstance(search, list), "The type of search is not List"

        out_dict = []
        for i in range(len(search)):
            # 传递gamma_scale参数
            x, aux_dict = self.backbone(
                z=template.copy(),
                x=search[i],
                l=list(descript[i]),
                temporal_query=self.track_query,
                top_K=self.token_len,
                gamma_scale=gamma_scale  # 添加参数
            )

            # 提取当前帧的重要tokens
            current_tokens = aux_dict["temproal_token"]  # [B, K, C]

            # 保存到token bank并维护FIFO队列
            self.token_bank.append(current_tokens.detach())
            if len(self.token_bank) > self.token_bank_max_len:
                self.token_bank.pop(0)

            # 处理历史token信息(如果有)
            if len(self.token_bank) > 1:
                # 拼接所有历史tokens
                all_bank = torch.cat(self.token_bank, dim=1)  # [B, N*K, C]
                B, _, C = all_bank.shape

                # 使用可学习Query融合历史tokens
                query = self.fusion_query.expand(B, -1, -1)  # [B, K, C]
                fused_out, _ = self.fusion_attn(query, all_bank, all_bank)
                fused_out = self.norm1(fused_out + query)  # 残差连接
                fused_out = self.norm2(fused_out + self.fusion_mlp(fused_out))

                # 与当前搜索区域交互(可选)
                feat_last = x
                if isinstance(x, list):
                    feat_last = x[-1]

                search_feat = feat_last[:, -self.feat_len_s:]  # 搜索区域特征
                new_out, _ = self.search_rebuild_attn(fused_out, search_feat, search_feat)
                new_out = new_out + fused_out  # 残差连接

                # 更新track_query
                self.track_query = new_out.detach()
            else:
                # 第一帧，直接使用当前tokens
                self.track_query = current_tokens.detach()

            feat_last = x
            if isinstance(x, list):
                feat_last = x[-1]

            enc_opt = feat_last[:, -self.feat_len_s:]  # encoder output for the search region (B, HW, C)

            att = torch.matmul(enc_opt, x[:, :1].transpose(1, 2))  # (B, HW, N)
            opt = (enc_opt.unsqueeze(-1) * att.unsqueeze(-2)).permute(
                (0, 3, 2, 1)).contiguous()  # (B, HW, C, N) --> (B, N, C, HW)

            # Forward head
            out = self.forward_head(opt, None)

            out.update(aux_dict)
            out['backbone_feat'] = x

            out_dict.append(out)

        return out_dict

    def forward_head(self, opt, gt_score_map=None):
        """
        enc_opt: output embeddings of the backbone, it can be (HW1+HW2, B, C) or (HW2, B, C)
        """
        # opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
        bs, Nq, C, HW = opt.size()
        opt_feat = opt.view(-1, C, self.feat_sz_s, self.feat_sz_s)

        if self.head_type == "CORNER":
            # run the corner head
            pred_box, score_map = self.box_head(opt_feat, True)
            outputs_coord = box_xyxy_to_cxcywh(pred_box)
            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
            out = {'pred_boxes': outputs_coord_new,
                   'score_map': score_map,
                   }
            return out

        elif self.head_type == "CENTER":
            # run the center head
            score_map_ctr, bbox, size_map, offset_map = self.box_head(opt_feat, gt_score_map)

            # outputs_coord = box_xyxy_to_cxcywh(bbox)
            outputs_coord = bbox
            outputs_coord_new = outputs_coord.view(bs, Nq, 4)

            out = {'pred_boxes': outputs_coord_new,
                   'score_map': score_map_ctr,
                   'size_map': size_map,
                   'offset_map': offset_map}

            return out
        else:
            raise NotImplementedError


def build_dutrack(cfg, training=True):
    current_dir = os.path.dirname(os.path.abspath(__file__))  # This is your Project Root
    pretrained_path = os.path.join(current_dir, '../../../pretrained_models')

    if cfg.MODEL.PRETRAIN_FILE and ('OSTrack' not in cfg.MODEL.PRETRAIN_FILE) and training:
        pretrained = os.path.join(pretrained_path, cfg.MODEL.PRETRAIN_FILE)
    else:
        pretrained = ''

    if cfg.MODEL.BACKBONE.TYPE == 'itpn_base':
        backbone = fast_itpn_base_3324_patch16_224(pretrained, drop_path_rate=cfg.TRAIN.DROP_PATH_RATE,
                                                   bert_dir=cfg.MODEL.BACKBONE.BERT_DIR)
    else:
        raise NotImplementedError

    hidden_dim = backbone.embed_dim
    patch_start_index = 1

    backbone.finetune_track(cfg=cfg, patch_start_index=patch_start_index)

    box_head = build_box_head(cfg, hidden_dim)

    model = DUTrack(
        backbone,
        box_head,
        aux_loss=False,
        head_type=cfg.MODEL.HEAD.TYPE,
        token_len=cfg.MODEL.BACKBONE.TOP_K,
    )
    if 'DUTrack' in cfg.MODEL.PRETRAIN_FILE and training:
        current_dir = os.path.dirname(os.path.abspath(__file__))  # This is your Project Root
        pretrained_path = os.path.join(current_dir, '../../../pretrained_models')
        file_name = cfg.MODEL.PRETRAIN_FILE
        pth = os.path.join(pretrained_path, file_name)
        checkpoint = torch.load(pth, map_location="cpu")
        missing_keys, unexpected_keys = model.load_state_dict(checkpoint["net"], strict=False)
        print('Load pretrained model from: ' + cfg.MODEL.PRETRAIN_FILE)
    return model
